import torch
import torch.nn as nn
import torch.optim as optim
from torchvision import datasets, transforms
import numpy as np

# 定义最简单线性模型
class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.fc = nn.Linear(28*28, 10)
    def forward(self, x):
        x = x.view(-1, 28*28)
        return self.fc(x)

# 训练模型并导出参数
def main():
    # 数据加载
    transform = transforms.Compose([transforms.ToTensor()])
    trainset = datasets.MNIST('.', train=True, download=True, transform=transform)
    loader = torch.utils.data.DataLoader(trainset, batch_size=64, shuffle=True)

    # 初始化模型、损失与优化器
    net = Net()
    criterion = nn.CrossEntropyLoss()
    optimizer = optim.SGD(net.parameters(), lr=0.1)

    # 简单训练5个epoch
    for epoch in range(5):
        for data, label in loader:
            optimizer.zero_grad()
            output = net(data)
            loss = criterion(output, label)
            loss.backward()
            optimizer.step()
        print(f"Epoch {epoch+1} 完成")

    # 获取 NumPy 参数
    W = net.fc.weight.data.numpy()
    b = net.fc.bias.data.numpy()
    # 量化到 int8
    scale = np.max(np.abs(W)) / 127
    Wq = np.round(W / scale).astype(np.int8)
    bq = np.round(b / scale).astype(np.int8)

    # 写入 mnist_params.h
    with open('mnist_params.h', 'w') as f:
        f.write('#ifndef MNIST_PARAMS_H\n')
        f.write('#define MNIST_PARAMS_H\n')
        f.write('#include <stdint.h>\n')
        f.write('#define IMAGE_SIZE (28*28)\n')
        f.write('extern const int8_t W[10][IMAGE_SIZE];\n')
        f.write('extern const int8_t b[10];\n')
        f.write('#endif')

    # 写入 mnist_params.c
    with open('mnist_params.c', 'w') as f:
        f.write('#include "mnist_params.h"\n')
        # 展开权重
        f.write('const int8_t W[10][IMAGE_SIZE] = {\n')
        for k in range(10):
            row = ','.join(str(int(x)) for x in Wq[k])
            f.write(f'    {{ {row} }},\n')
        f.write('};\n')
        # 展开偏置
        b_row = ','.join(str(int(x)) for x in bq)
        f.write(f'const int8_t b[10] = {{ {b_row} }};\n')

    print('参数已导出到 mnist_params.h 和 mnist_params.c')

if __name__ == '__main__':
    main()